# coding=utf-8
# 格式转换
# MindSpore中可以把用于训练网络模型的数据集，转换为MindSpore特定的格式数据（MindSpore Record格式），从而更加方便地保存和加载数据
# MindSpore Record数据格式具备的特征如下：
# 实现数据统一存储、访问，使得训练时数据读取更加简便。
# 数据聚合存储、高效读取，使得训练时数据方便管理和移动。
# 高效的数据编解码操作，使得用户可以对数据操作无感知。
# 可以灵活控制数据切分的分区大小，实现分布式数据处理。

# 转换成Record格式
from PIL import Image
from io import BytesIO
from mindspore.mindrecord import FileWriter

file_name = "test_vision.mindrecord"
# 定义包含的字段
cv_schema = {"file_name": {"type": "string"},
             "label": {"type": "int32"},
             "data": {"type": "bytes"}}

# 声明MindSpore Record文件格式
writer = FileWriter(file_name, shard_num=1, overwrite=True)
writer.add_schema(cv_schema, "it is a cv dataset")
writer.add_index(["file_name", "label"])

# 创建数据集
data = []
for i in range(100):
    sample = {}
    white_io = BytesIO()
    Image.new('RGB', ((i + 1) * 10, (i + 1) * 10), (255, 255, 255)).save(white_io, 'JPEG')
    image_bytes = white_io.getvalue()
    sample['file_name'] = str(i + 1) + ".jpg"
    sample['label'] = i + 1
    sample['data'] = white_io.getvalue()

    data.append(sample)
    if i % 10 == 0:
        writer.write_raw_data(data)
        data = []

if data:
    writer.write_raw_data(data)

writer.commit()

# 通过MindDataset接口读取MindSpore Record文件格式。
from mindspore.dataset import MindDataset
from mindspore.dataset.vision import Decode

# 读取MindSpore Record文件格式
data_set = MindDataset(dataset_files=file_name)
decode_op = Decode()
data_set = data_set.map(operations=decode_op, input_columns=["data"], num_parallel_workers=2)

# 样本计数
print("Got {} samples".format(data_set.get_dataset_size()))

# 转换NLP类数据集
# 1、生成100条文本数据，并转换成MindSpore Record文件格式。
import numpy as np
from mindspore.mindrecord import FileWriter

# 输出的MindSpore Record文件完整路径
file_name = "test_text.mindrecord"

# 定义样本数据包含的字段
nlp_schema = {"source_sos_ids": {"type": "int64", "shape": [-1]},
              "source_sos_mask": {"type": "int64", "shape": [-1]},
              "source_eos_ids": {"type": "int64", "shape": [-1]},
              "source_eos_mask": {"type": "int64", "shape": [-1]},
              "target_sos_ids": {"type": "int64", "shape": [-1]},
              "target_sos_mask": {"type": "int64", "shape": [-1]},
              "target_eos_ids": {"type": "int64", "shape": [-1]},
              "target_eos_mask": {"type": "int64", "shape": [-1]}}

# 声明MindSpore Record文件格式
writer = FileWriter(file_name, shard_num=1, overwrite=True)
writer.add_schema(nlp_schema, "Preprocessed nlp dataset.")

# 创建虚拟数据集
data = []
for i in range(100):
    sample = {"source_sos_ids": np.array([i, i + 1, i + 2, i + 3, i + 4], dtype=np.int64),
              "source_sos_mask": np.array([i * 1, i * 2, i * 3, i * 4, i * 5, i * 6, i * 7], dtype=np.int64),
              "source_eos_ids": np.array([i + 5, i + 6, i + 7, i + 8, i + 9, i + 10], dtype=np.int64),
              "source_eos_mask": np.array([19, 20, 21, 22, 23, 24, 25, 26, 27], dtype=np.int64),
              "target_sos_ids": np.array([28, 29, 30, 31, 32], dtype=np.int64),
              "target_sos_mask": np.array([33, 34, 35, 36, 37, 38], dtype=np.int64),
              "target_eos_ids": np.array([39, 40, 41, 42, 43, 44, 45, 46, 47], dtype=np.int64),
              "target_eos_mask": np.array([48, 49, 50, 51], dtype=np.int64)}
    data.append(sample)

    if i % 10 == 0:
        writer.write_raw_data(data)
        data = []

if data:
    writer.write_raw_data(data)

writer.commit()

# 读取MindSpore Record文件格式
data_set = MindDataset(dataset_files=file_name, shuffle=False)

# 样本计数
print("Got {} samples".format(data_set.get_dataset_size()))

# 打印部分数据
count = 0
for item in data_set.create_dict_iterator(output_numpy=True):
    print("source_sos_ids:", item["source_sos_ids"])
    count += 1
    if count == 10:
        break

# Dataset转存MindRecord
# 转存CIFAR-10数据集
# from download import download
from mindspore.dataset import Cifar10Dataset

# url = "https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/datasets/cifar-10-binary.tar.gz"
# path = download(url, "./", kind="tar.gz", replace=True)
local_path = "../../datasets/cifar-10-batches-bin/"
mind_record = local_path + "cifar10.mindrecord"
dataset = Cifar10Dataset(local_path)  # 加载数据
dataset.save(mind_record)
import os
from mindspore.dataset import MindDataset

# 读取MindSpore Record文件格式
data_set = MindDataset(dataset_files=mind_record)

# 样本计数
print("Got {} samples".format(data_set.get_dataset_size()))

if os.path.exists(mind_record) and os.path.exists(local_path + "cifar10.mindrecord.db"):
    os.remove(mind_record)
    os.remove(local_path + "cifar10.mindrecord.db")
